from typing import Dict, Any, List
from agents.components.base import BaseComponent, ComponentConfig, ComponentType

class EmbeddingComponent(BaseComponent):
    """文本向量化组件的基类"""
    
    @classmethod
    def get_component_type(cls) -> ComponentType:
        return ComponentType.EMBEDDING
    
    @classmethod
    def get_component_name(cls) -> str:
        raise NotImplementedError("子类必须实现 get_component_name 方法")
    
    @classmethod
    def get_component_description(cls) -> str:
        raise NotImplementedError("子类必须实现 get_component_description 方法")
    
    @classmethod
    def get_config_schema(cls) -> Dict[str, Any]:
        base_schema = super().get_config_schema()
        base_schema["properties"].update({
            "model_name": {
                "type": "string",
                "description": "向量化模型名称"
            }
        })
        return base_schema
    
    @classmethod
    def get_input_schema(cls) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "texts": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "需要向量化的文本列表"
                }
            },
            "required": ["texts"]
        }
    
    @classmethod
    def get_output_schema(cls) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "embeddings": {
                    "type": "array",
                    "items": {
                        "type": "array",
                        "items": {"type": "number"}
                    },
                    "description": "文本的向量表示列表"
                }
            },
            "required": ["embeddings"]
        }
    
    def __init__(self, config: ComponentConfig):
        super().__init__(config)
        self.model_name = config.config.get("model_name")
        
    def validate_input(self, input_data: Dict[str, Any]) -> bool:
        if not isinstance(input_data, dict):
            return False
        if "texts" not in input_data:
            return False
        if not isinstance(input_data["texts"], list):
            return False
        return all(isinstance(text, str) for text in input_data["texts"])
    
    def validate_output(self, output_data: Dict[str, Any]) -> bool:
        if not isinstance(output_data, dict):
            return False
        if "embeddings" not in output_data:
            return False
        if not isinstance(output_data["embeddings"], list):
            return False
        return all(isinstance(embedding, list) and 
                  all(isinstance(x, (int, float)) for x in embedding) 
                  for embedding in output_data["embeddings"])
    
    async def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """执行文本向量化
        
        Args:
            input_data: 包含以下字段的字典
                - texts: List[str], 需要向量化的文本列表
                
        Returns:
            Dict[str, Any]: 包含以下字段的字典
                - embeddings: List[List[float]], 文本的向量表示列表
        """
        if not self.validate_input(input_data):
            raise ValueError("输入数据验证失败")
            
        if "texts" not in input_data:
            raise ValueError("输入数据必须包含 texts 字段")
            
        if not isinstance(input_data["texts"], list):
            raise ValueError("texts 必须是列表类型")
            
        return await self._embed_texts(input_data["texts"])
    
    async def _embed_texts(self, texts: List[str]) -> Dict[str, Any]:
        """具体的向量化实现，由子类实现
        
        Args:
            texts: 需要向量化的文本列表
            
        Returns:
            Dict[str, Any]: 包含 embeddings 字段的字典
        """
        raise NotImplementedError("子类必须实现 _embed_texts 方法") 